package com.codeanalysissys.backend.golang;

import com.codeanalysissys.backend.javascript.parser.JavaScriptLexer;
import com.codeanalysissys.backend.javascript.parser.JavaScriptParser;
import ghaffarian.graphs.*;

import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.nio.file.Paths;
import java.util.*;

import com.codeanalysissys.backend.golang.parser.GoLexer;
import com.codeanalysissys.backend.golang.parser.GoParser;
import org.antlr.v4.runtime.*;
import org.antlr.v4.runtime.misc.Interval;
import org.antlr.v4.runtime.tree.ParseTree;
import com.codeanalysissys.backend.graphs.cfg.CFEdge;
import com.codeanalysissys.backend.graphs.cfg.CFNode;
import com.codeanalysissys.backend.graphs.cfg.ControlFlowGraph;
import com.codeanalysissys.backend.golang.parser.GoParserBaseVisitor;
import ghaffarian.nanologger.Logger;

/**
 * A Control Flow Graph (CFG) builder for Go programs.
 * A Go parser generated via ANTLRv4 is used for this purpose.
 * This implementation is based on ANTLRv4's Visitor pattern.
 *
 * @author Huang WeiQuan
 */
public class GoCFGBuilder {
    /**
     * Build and return the Control Flow Graph (CFG) for the given Go source file.
     */

    public static ControlFlowGraph build(String goFile) throws IOException {
        return build(new File(goFile));
    }

    public static ControlFlowGraph build(InputStream inFile) throws IOException {
        // 将字符串转换为字节数组
//        byte[] bytes = content.getBytes();
        // 使用 ByteArrayInputStream 封装字节数组为输入流
//        InputStream inFile = new ByteArrayInputStream(bytes);

        ANTLRInputStream input = new ANTLRInputStream(inFile);
        GoLexer lexer = new GoLexer(input);
        CommonTokenStream tokens = new CommonTokenStream(lexer);
        GoParser parser = new GoParser(tokens);
        ParseTree tree = parser.sourceFile();
        return build(tree, null, null);
    }

    /**
     * Build and return the Control Flow Graph (CFG) for the given Go source file.
     */

    public static ControlFlowGraph build(File goFile) throws IOException {
        if (!goFile.getName().endsWith(".go"))
            throw new IOException("Not a Golang File!");
        CharStream input = CharStreams.fromPath(Paths.get(goFile.getAbsolutePath()));
        GoLexer lexer = new GoLexer(input);
        CommonTokenStream tokens = new CommonTokenStream(lexer);
        GoParser parser = new GoParser(tokens);
        ParseTree tree = parser.sourceFile();
        return build(goFile.getName(), tree, null, null);
    }

    public static ControlFlowGraph build(String filename, InputStream inputStream) throws IOException {
        InputStream inFile = inputStream;
        ANTLRInputStream input = new ANTLRInputStream(inFile);
        GoLexer lexer = new GoLexer(input);
        CommonTokenStream tokens = new CommonTokenStream(lexer);
        GoParser parser = new GoParser(tokens);
        ParseTree tree = parser.sourceFile();                               //Java语法规则中的根规则-解析器已生成
        return build(filename, tree, null, null);
    }

    /**
     * Build and return the Control Flow Graph (CFG) for the given Parse-Tree.
     * The 'ctxProps' map includes contextual-properties for particular nodes
     * in the parse-tree, which can be used for linking this graph with other
     * graphs by using the same parse-tree and the same contextual-properties.
     */

    public static ControlFlowGraph build(ParseTree tree, String propKey, Map<ParserRuleContext, Object> ctxProps) {
        ControlFlowGraph cfg = new ControlFlowGraph();
        ControlFlowVisitor visitor = new ControlFlowVisitor(cfg, propKey, ctxProps);
        visitor.visit(tree);
        return cfg;
    }

    public static ControlFlowGraph build(String goFileName, ParseTree tree,
                                         String propKey, Map<ParserRuleContext, Object> ctxProps) {
        ControlFlowGraph cfg = new ControlFlowGraph(goFileName);
        ControlFlowVisitor visitor = new ControlFlowVisitor(cfg, propKey, ctxProps);
        visitor.visit(tree);
        visitor.handleDefer();
        return cfg;
    }

    public static class ControlFlowVisitor extends GoParserBaseVisitor<Void> {
        private ControlFlowGraph cfg;
        private Deque<CFNode> preNodes;
        private Deque<CFEdge.Type> preEdges;
        private Deque<Block> loopBlocks;
        private List<Block> labeledBlocks;
        private Deque<Block> tryBlocks;
        private Queue<CFNode> casesQueue;
        private List<CFNode> DeferNode;
        private List<GoParser.DeferStmtContext> DeferState;
        private boolean dontPop;
        private String propKey;
        private Map<ParserRuleContext, Object> contexutalProperties;
        private int numDefer = 0;

        /*private Deque<String> classNames;*/
        public ControlFlowVisitor(ControlFlowGraph cfg, String propKey, Map<ParserRuleContext, Object> ctxProps) {
            preNodes = new ArrayDeque<>();
            preEdges = new ArrayDeque<>();
            loopBlocks = new ArrayDeque<>();
            labeledBlocks = new ArrayList<>();
            tryBlocks = new ArrayDeque<>();
            casesQueue = new ArrayDeque<>();
            DeferNode = new ArrayList<>();
            DeferState = new ArrayList<>();
            /*classNames = new ArrayDeque<>();*/
            dontPop = false;
            this.cfg = cfg;
            this.propKey = propKey;
            contexutalProperties = ctxProps;
        }

        /**
         * Reset all data-structures and flags for visiting a new method declaration.
         */
        private void init() {
            preNodes.clear();
            preEdges.clear();
            loopBlocks.clear();
            labeledBlocks.clear();
            tryBlocks.clear();
            dontPop = false;
        }

        /**
         * Add contextual properties to the given node.
         * This will first check to see if there is any property for the
         * given context, and if so, the property will be added to the node.
         */

        private void addContextualProperty(CFNode node, ParserRuleContext ctx) {
            if (propKey != null && contexutalProperties != null) {
                Object prop = contexutalProperties.get(ctx);
                if (prop != null)
                    node.setProperty(propKey, prop);
            }
        }

        @Override
        public Void visitPackageClause(GoParser.PackageClauseContext ctx) {
            // packageClause :  annotation* 'package' qualifiedName ';'
            if (ctx != null && ctx.IDENTIFIER() != null)
                cfg.setPackage(ctx.IDENTIFIER().getText());
            return null;
        }

        /**
         * Golang doesn't have enum and class.
         * So this CFGBuilder will not visit enum, class,constructor and class body.
         */

        @Override
        public Void visitInterfaceType(GoParser.InterfaceTypeContext ctx) {
            //Just ignore the interface for the moment...
            return null;
        }

        @Override
        public Void visitFunctionDecl(GoParser.FunctionDeclContext ctx) {
            init();
            CFNode entry = new CFNode();
            entry.setLineOfCode(ctx.getStart().getLine());
            String retType = "func";
            String args = getOriginalCodeText(ctx.signature());
            entry.setCode(retType + " " + ctx.IDENTIFIER() + args);
            addContextualProperty(entry, ctx);
            cfg.addVertex(entry);
            entry.setProperty("name", ctx.IDENTIFIER().getText());
            entry.setProperty("type", retType);
            cfg.addMethodEntry(entry);
            preNodes.push(entry);
            preEdges.push(CFEdge.Type.EPSILON);
            return visitChildren(ctx);
        }

        @Override
        public Void visitExpressionStmt(GoParser.ExpressionStmtContext ctx) {
            CFNode expr = new CFNode();
            expr.setLineOfCode(ctx.getStart().getLine());
            expr.setCode(getOriginalCodeText(ctx));
            Logger.debug(expr.getLineOfCode() + ": " + expr.getCode());
            addContextualProperty(expr, ctx);
            addNodeAndPreEdge(expr);
            preEdges.push(CFEdge.Type.EPSILON);
            preNodes.push(expr);
            return null;
        }

        @Override
        public Void visitVarDecl(GoParser.VarDeclContext ctx) {
            CFNode declr = new CFNode();
            declr.setLineOfCode(ctx.getStart().getLine());
            declr.setCode(getOriginalCodeText(ctx));
            addContextualProperty(declr, ctx);
            addNodeAndPreEdge(declr);
            preEdges.push(CFEdge.Type.EPSILON);
            preNodes.push(declr);
            return null;
        }

        @Override
        public Void visitSimpleStmt(GoParser.SimpleStmtContext ctx) {
            if (ctx.shortVarDecl() != null)
                return visit(ctx.shortVarDecl());
            CFNode declr = new CFNode();
            declr.setLineOfCode(ctx.getStart().getLine());
            declr.setCode(getOriginalCodeText(ctx));
            addContextualProperty(declr, ctx);
            addNodeAndPreEdge(declr);
            preEdges.push(CFEdge.Type.EPSILON);
            preNodes.push(declr);
            return null;
        }

        @Override
        public Void visitIfStmt(GoParser.IfStmtContext ctx) {
            CFNode ifNode = new CFNode();
            ifNode.setLineOfCode(ctx.getStart().getLine());
            ifNode.setCode("if " + getOriginalCodeText(ctx.expression()));
            addContextualProperty(ifNode, ctx);
            addNodeAndPreEdge(ifNode);
            preEdges.push(CFEdge.Type.TRUE);
            preNodes.push(ifNode);
            visit(ctx.block(0));
            CFNode endif = new CFNode();
            endif.setLineOfCode(0);
            endif.setCode("endif");
            addNodeAndPreEdge(endif);
            if (ctx.block().size() == 1) { // if without else
                cfg.addEdge(new Edge<>(ifNode, new CFEdge(CFEdge.Type.FALSE), endif));
            } else {  //  if with else
                preEdges.push(CFEdge.Type.FALSE);
                preNodes.push(ifNode);
                visit(ctx.block(1));
                popAddPreEdgeTo(endif);
            }
            preEdges.push(CFEdge.Type.EPSILON);
            preNodes.push(endif);
            return null;
        }

        @Override
        public Void visitShortVarDecl(GoParser.ShortVarDeclContext ctx) {
            CFNode declr = new CFNode();
            declr.setLineOfCode(ctx.getStart().getLine());
            declr.setCode(getOriginalCodeText(ctx));
            addContextualProperty(declr, ctx);
            addNodeAndPreEdge(declr);
            preEdges.push(CFEdge.Type.EPSILON);
            preNodes.push(declr);
            return null;
        }

        @Override
        public Void visitForStmt(GoParser.ForStmtContext ctx) {
            if (ctx.forClause() != null) //Common for loop
            {
                CFNode forInit = null;
                if (ctx.forClause().initStmt != null) {
                    forInit = new CFNode();
                    forInit.setLineOfCode(ctx.forClause().initStmt.getStart().getLine());
                    forInit.setCode(getOriginalCodeText(ctx.forClause().initStmt));
                    addContextualProperty(forInit, ctx.forClause().initStmt);
                    addNodeAndPreEdge(forInit);
                }
                CFNode forExpr = new CFNode();
                if (ctx.forClause().expression() == null) {
                    forExpr.setLineOfCode(ctx.forClause().getStart().getLine());
                    forExpr.setCode("for ; ");
                } else {
                    forExpr.setLineOfCode(ctx.forClause().expression().getStart().getLine());
                    forExpr.setCode("for " + getOriginalCodeText(ctx.forClause().expression()) + " ");
                }
                addContextualProperty(forExpr, ctx.forClause().expression());
                cfg.addVertex(forExpr);
                if (forInit != null)
                    cfg.addEdge(new Edge<>(forInit, new CFEdge(CFEdge.Type.EPSILON), forExpr));
                else
                    popAddPreEdgeTo(forExpr);
                CFNode forUpdate = new CFNode();
                if (ctx.forClause().postStmt == null) {
                    forUpdate.setCode(" ; ");
                    forUpdate.setLineOfCode(ctx.forClause().getStart().getLine());
                } else {
                    forUpdate.setCode(getOriginalCodeText(ctx.forClause().postStmt));
                    forUpdate.setLineOfCode(ctx.forClause().postStmt.getStart().getLine());
                }
                addContextualProperty(forUpdate, ctx.forClause().postStmt);
                cfg.addVertex(forUpdate);
                CFNode forEnd = new CFNode();
                forEnd.setLineOfCode(0);
                forEnd.setCode("endfor");
                cfg.addVertex(forEnd);
                cfg.addEdge(new Edge<>(forExpr, new CFEdge(CFEdge.Type.FALSE), forEnd));
                preEdges.push(CFEdge.Type.TRUE);
                preNodes.push(forExpr);
                loopBlocks.push(new Block(forUpdate, forEnd));
                visit(ctx.block());
                loopBlocks.pop();
                popAddPreEdgeTo(forUpdate);
                cfg.addEdge(new Edge<>(forUpdate, new CFEdge(CFEdge.Type.EPSILON), forExpr));
                preEdges.push(CFEdge.Type.EPSILON);
                preNodes.push(forEnd);
            } else {
                CFNode forExpr = new CFNode();
                if (ctx.expression() != null) {
                    forExpr.setLineOfCode(ctx.expression().getStart().getLine());
                    forExpr.setCode("for " + getOriginalCodeText(ctx.expression()) + "");
                } else {
                    forExpr.setLineOfCode(0);
                    forExpr.setCode("for" + "");
                }
                addContextualProperty(forExpr, ctx.expression());
                addNodeAndPreEdge(forExpr);
                CFNode forEnd = new CFNode();
                forEnd.setLineOfCode(0);
                forEnd.setCode("endfor");
                cfg.addVertex(forEnd);
                cfg.addEdge(new Edge<>(forExpr, new CFEdge(CFEdge.Type.FALSE), forEnd));
                //
                preEdges.push(CFEdge.Type.TRUE);
                preNodes.push(forExpr);
                //
                loopBlocks.push(new Block(forExpr, forEnd));
                visit(ctx.block());
                loopBlocks.pop();
                popAddPreEdgeTo(forExpr);
                //
                preEdges.push(CFEdge.Type.EPSILON);
                preNodes.push(forEnd);
            }
            return null;
        }

        @Override
        public Void visitSwitchStmt(GoParser.SwitchStmtContext ctx) {
            CFNode switchNode = new CFNode();
            switchNode.setLineOfCode(ctx.getStart().getLine());
            if (ctx.exprSwitchStmt().expression() != null) {
                switchNode.setCode("switch " + getOriginalCodeText(ctx.exprSwitchStmt().expression()));
            } else {
                switchNode.setCode("switch ");
            }
            addContextualProperty(switchNode, ctx);
            addNodeAndPreEdge(switchNode);
            CFNode endSwitch = new CFNode();
            endSwitch.setLineOfCode(0);
            endSwitch.setCode("end-switch");
            cfg.addVertex(endSwitch);
            preEdges.push(CFEdge.Type.EPSILON);
            preNodes.push(switchNode);
            loopBlocks.push(new Block(switchNode, endSwitch));
            CFNode preCase = null;
            for (GoParser.ExprCaseClauseContext grp : ctx.exprSwitchStmt().exprCaseClause()) {
                preCase = visitexprSwitchCase(grp.exprSwitchCase(), preCase);
                for (GoParser.StatementListContext blk : grp.statementList())
                    visit(blk);
            }
            preCase = visitexprSwitchCase(ctx.exprSwitchStmt().exprCaseClause(0).exprSwitchCase(), preCase);
            loopBlocks.pop();
            popAddPreEdgeTo(endSwitch);
            if (preCase != null)
                cfg.addEdge(new Edge<>(preCase, new CFEdge(CFEdge.Type.FALSE), endSwitch));
            preEdges.push(CFEdge.Type.EPSILON);
            preNodes.push(endSwitch);
            return null;
        }

        private CFNode visitexprSwitchCase(List<GoParser.ExprSwitchCaseContext> list, CFNode preCase) {
            CFNode caseStmnt = preCase;
            for (GoParser.ExprSwitchCaseContext ctx : list) {
                caseStmnt = new CFNode();
                caseStmnt.setLineOfCode(ctx.getStart().getLine());
                caseStmnt.setCode(getOriginalCodeText(ctx));
                cfg.addVertex(caseStmnt);
                if (dontPop)
                    dontPop = false;
                else
                    cfg.addEdge(new Edge<>(preNodes.pop(), new CFEdge(preEdges.pop()), caseStmnt));
                if (preCase != null)
                    cfg.addEdge(new Edge<>(preCase, new CFEdge(CFEdge.Type.FALSE), caseStmnt));
                if (ctx.getStart().getText().equals("default")) {
                    preEdges.push(CFEdge.Type.EPSILON);
                    preNodes.push(caseStmnt);
                    caseStmnt = null;
                } else { // any other case ...
                    dontPop = true;
                    casesQueue.add(caseStmnt);
                    preCase = caseStmnt;
                }
            }
            return caseStmnt;
        }

        @Override
        public Void visitLabeledStmt(GoParser.LabeledStmtContext ctx) {
            CFNode labelNode = new CFNode();
            labelNode.setLineOfCode(ctx.getStart().getLine());
            labelNode.setCode(ctx.IDENTIFIER() + ": ");
            addContextualProperty(labelNode, ctx);
            addNodeAndPreEdge(labelNode);
            CFNode endLabelNode = new CFNode();
            endLabelNode.setLineOfCode(0);
            endLabelNode.setCode("end-label");
            cfg.addVertex(endLabelNode);
            preEdges.push(CFEdge.Type.EPSILON);
            preNodes.push(labelNode);
            labeledBlocks.add(new Block(labelNode, endLabelNode, ctx.IDENTIFIER().getText()));
            visit(ctx.statement());
            popAddPreEdgeTo(endLabelNode);
            preEdges.push(CFEdge.Type.EPSILON);
            preNodes.push(endLabelNode);
            return null;
        }

        @Override
        public Void visitReturnStmt(GoParser.ReturnStmtContext ctx) {
            CFNode ret = new CFNode();
            ret.setLineOfCode(ctx.getStart().getLine());
            ret.setCode(getOriginalCodeText(ctx));
            addContextualProperty(ret, ctx);
            addNodeAndPreEdge(ret);
            dontPop = true;
            return null;
        }

        @Override
        public Void visitBreakStmt(GoParser.BreakStmtContext ctx) {
            CFNode breakNode = new CFNode();
            breakNode.setLineOfCode(ctx.getStart().getLine());
            breakNode.setCode(getOriginalCodeText(ctx));
            addContextualProperty(breakNode, ctx);
            addNodeAndPreEdge(breakNode);
            if (ctx.IDENTIFIER() != null) {
                // a label is specified
                for (Block block : labeledBlocks) {
                    if (block.label.equals(ctx.IDENTIFIER().getText())) {
                        cfg.addEdge(new Edge<>(breakNode, new CFEdge(CFEdge.Type.EPSILON), block.end));
                        break;
                    }
                }
            } else {
                // no label
                Block block = loopBlocks.peek();
                cfg.addEdge(new Edge<>(breakNode, new CFEdge(CFEdge.Type.EPSILON), block.end));
            }
            dontPop = true;
            return null;
        }

        @Override
        public Void visitContinueStmt(GoParser.ContinueStmtContext ctx) {
            CFNode continueNode = new CFNode();
            continueNode.setLineOfCode(ctx.getStart().getLine());
            continueNode.setCode(getOriginalCodeText(ctx));
            addContextualProperty(continueNode, ctx);
            addNodeAndPreEdge(continueNode);
            if (ctx.IDENTIFIER() != null) {
                // a label is specified
                for (Block block : labeledBlocks) {
                    if (block.label.equals(ctx.IDENTIFIER().getText())) {
                        cfg.addEdge(new Edge<>(continueNode, new CFEdge(CFEdge.Type.EPSILON), block.start));
                        break;
                    }
                }
            } else {
                // no label
                Block block = loopBlocks.peek();
                cfg.addEdge(new Edge<>(continueNode, new CFEdge(CFEdge.Type.EPSILON), block.start));
            }
            dontPop = true;
            return null;
        }

        @Override
        public Void visitEmptyStmt(GoParser.EmptyStmtContext ctx) {
            //Empty
            return null;
        }

        @Override
        public Void visitIncDecStmt(GoParser.IncDecStmtContext ctx) {
            CFNode incDec = new CFNode();
            incDec.setLineOfCode(ctx.getStart().getLine());
            incDec.setCode(getOriginalCodeText(ctx));
            addContextualProperty(incDec, ctx);
            addNodeAndPreEdge(incDec);
            preEdges.push(CFEdge.Type.EPSILON);
            preNodes.push(incDec);
            return null;
        }

        @Override
        public Void visitGotoStmt(GoParser.GotoStmtContext ctx) {
            CFNode gotoNode = new CFNode();
            gotoNode.setLineOfCode(ctx.getStart().getLine());
            gotoNode.setCode(getOriginalCodeText(ctx));
            addContextualProperty(gotoNode, ctx);
            addNodeAndPreEdge(gotoNode);
            int temp = 0;
            if (ctx.IDENTIFIER() != null) {
                // a label is specified
                for (Block block : labeledBlocks) {
                    if (block.label.equals(ctx.IDENTIFIER().getText())) {
                        cfg.addEdge(new Edge<>(gotoNode, new CFEdge(CFEdge.Type.EPSILON), block.start));
                        temp = 1;
                        break;
                    }
                }
                if (temp == 0) {
                    preEdges.push(CFEdge.Type.EPSILON);
                    preNodes.push(gotoNode);
                }
                dontPop = true;
            }
            return null;
        }

        @Override
        public Void visitFallthroughStmt(GoParser.FallthroughStmtContext ctx) {
            CFNode fallthroughNode = new CFNode();
            fallthroughNode.setLineOfCode(ctx.getStart().getLine());
            fallthroughNode.setCode(getOriginalCodeText(ctx));
            addContextualProperty(fallthroughNode, ctx);
            addNodeAndPreEdge(fallthroughNode);
            CFNode nextLabel = new CFNode();
            nextLabel.setLineOfCode(0);
            nextLabel.setCode("next-case");
            addNodeAndPreEdge(nextLabel);
            cfg.addEdge(new Edge<>(fallthroughNode, new CFEdge(CFEdge.Type.EPSILON), nextLabel));
            dontPop = true;
            return null;
        }

        @Override
        public Void visitDeferStmt(GoParser.DeferStmtContext ctx) {
            CFNode deferNode = new CFNode();
            deferNode.setLineOfCode(ctx.getStart().getLine());
            deferNode.setCode(getOriginalCodeText(ctx.expression()));
            addContextualProperty(deferNode, ctx);
            DeferState.add(ctx);
            DeferNode.add(deferNode);
            numDefer++;
            return null;
        }

        @Override
        public Void visitSendStmt(GoParser.SendStmtContext ctx) {
            CFNode sendNode = new CFNode();
            sendNode.setLineOfCode(ctx.getStart().getLine());
            sendNode.setCode(getOriginalCodeText(ctx.expression(0)));
            addContextualProperty(sendNode, ctx);
            addNodeAndPreEdge(sendNode);
            preEdges.push(CFEdge.Type.EPSILON);
            preNodes.push(sendNode);
            return null;
        }

        @Override
        public Void visitRecvStmt(GoParser.RecvStmtContext ctx) {
            CFNode receiveNode = new CFNode();
            receiveNode.setLineOfCode(ctx.getStart().getLine());
            receiveNode.setCode(getOriginalCodeText(ctx.expression()));
            addContextualProperty(receiveNode, ctx);
            addNodeAndPreEdge(receiveNode);
            preEdges.push(CFEdge.Type.EPSILON);
            preNodes.push(receiveNode);
            return null;
        }

        @Override
        public Void visitSelectStmt(GoParser.SelectStmtContext ctx) {
            CFNode selectNode = new CFNode();
            selectNode.setLineOfCode(ctx.getStart().getLine());
            selectNode.setCode("select");
            addContextualProperty(selectNode, ctx);
            addNodeAndPreEdge(selectNode);
            CFNode endSelect = new CFNode();
            endSelect.setLineOfCode(0);
            endSelect.setCode("end-select");
            cfg.addVertex(endSelect);
            preEdges.push(CFEdge.Type.EPSILON);
            preNodes.push(selectNode);
            loopBlocks.push(new Block(selectNode, endSelect));
            CFNode preCase = null;
            for (GoParser.CommClauseContext grp : ctx.commClause()) {
                preCase = visitSelectCase(grp.statementList().statement(), preCase);
                for (GoParser.StatementListContext blk : grp.statementList(0))
                    visit(blk);
            }
            preCase = visitSelectCase(ctx.commClause(0).statementList().statement(), preCase);
            loopBlocks.pop();
            popAddPreEdgeTo(endSelect);
            if (preCase != null)
                cfg.addEdge(new Edge<>(preCase, new CFEdge(CFEdge.Type.FALSE), endSelect));
            preEdges.push(CFEdge.Type.EPSILON);
            preNodes.push(endSelect);
            return null;
        }

        private CFNode visitSelectCase(List<GoParser.StatementContext> list, CFNode preCase) {
            CFNode caseStmnt = preCase;
            for (GoParser.StatementContext ctx : list) {
                caseStmnt = new CFNode();
                caseStmnt.setLineOfCode(ctx.getStart().getLine());
                caseStmnt.setCode(getOriginalCodeText(ctx));
                cfg.addVertex(caseStmnt);
                if (dontPop)
                    dontPop = false;
                else
                    cfg.addEdge(new Edge<>(preNodes.pop(), new CFEdge(preEdges.pop()), caseStmnt));
                if (preCase != null)
                    cfg.addEdge(new Edge<>(preCase, new CFEdge(CFEdge.Type.FALSE), caseStmnt));
                if (ctx.getStart().getText().equals("default")) {
                    preEdges.push(CFEdge.Type.EPSILON);
                    preNodes.push(caseStmnt);
                    caseStmnt = null;
                } else { // any other case ...
                    dontPop = true;
                    casesQueue.add(caseStmnt);
                    preCase = caseStmnt;
                }
            }
            return caseStmnt;
        }

        @Override
        public Void visitMethodDecl(GoParser.MethodDeclContext ctx) {
            init();
            CFNode entry = new CFNode();
            entry.setLineOfCode(ctx.getStart().getLine());
            String retType = "func ";
            String retype = "method";
            String args = getOriginalCodeText(ctx.signature());
            String args2 = getOriginalCodeText(ctx.receiver());
            entry.setCode(retType + args2 + " " + ctx.IDENTIFIER() + args);
            addContextualProperty(entry, ctx);
            cfg.addVertex(entry);
            entry.setProperty("name", ctx.IDENTIFIER().getText());
            entry.setProperty("type", retype);
            cfg.addMethodEntry(entry);
            preNodes.push(entry);
            preEdges.push(CFEdge.Type.EPSILON);
            return visitChildren(ctx);
        }

        /**
         * Get resulting Control-Flow-Graph of this CFG-Builder.
         */
        public ControlFlowGraph getCFG() {
            return cfg;
        }

        /**
         * Handle the defer statement.
         */
        public void handleDefer() {
            while (numDefer != 0) {
                int temp = numDefer - 1;
                CFNode deferNode = DeferNode.get(temp);
                addNodeAndPreEdge(deferNode);
                preEdges.push(CFEdge.Type.EPSILON);
                preNodes.push(deferNode);
                visit(DeferState.get(temp).expression());
                DeferNode.remove(temp);
                numDefer--;
            }
        }

        /**
         * Add this node to the CFG and create edge from pre-node to this node.
         */
        private void addNodeAndPreEdge(CFNode node) {
            cfg.addVertex(node);
            popAddPreEdgeTo(node);
        }

        /**
         * Add a new edge to the given node, by poping the edge-type of the stack.
         */
        private void popAddPreEdgeTo(CFNode node) {
            if (dontPop)
                dontPop = false;
            else {
                Logger.debug("\nPRE-NODES = " + preNodes.size());
                Logger.debug("PRE-EDGES = " + preEdges.size() + '\n');
                cfg.addEdge(new Edge<>(preNodes.pop(), new CFEdge(preEdges.pop()), node));
            }
            //
            for (int i = casesQueue.size(); i > 0; --i)
                cfg.addEdge(new Edge<>(casesQueue.remove(), new CFEdge(CFEdge.Type.TRUE), node));
        }

        /**
         * Get the original program text for the given parser-rule context.
         * This is required for preserving whitespaces.
         */
        private String getOriginalCodeText(ParserRuleContext ctx) {
            int start = ctx.start.getStartIndex();
            int stop = ctx.stop.getStopIndex();
            Interval interval = new Interval(start, stop);
            return ctx.start.getInputStream().getText(interval);
        }

        /**
         * A simple structure for holding the start, end, and label of code blocks.
         * These are used to handle 'break' and 'continue' statements.
         */
        private class Block {

            public final String label;
            public final CFNode start, end;

            Block(CFNode start, CFNode end, String label) {
                this.start = start;
                this.end = end;
                this.label = label;
            }

            Block(CFNode start, CFNode end) {
                this(start, end, "");
            }
        }
    }
}